#include "KhaosVulkanSDK.h"
#define VMA_IMPLEMENTATION
#include "KhaosVmaUsage.h"

#ifndef _WIN32
#define EXTERNAL_MEMORY_HANDLE_SUPPORTED_TYPE    VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR
#else
#define EXTERNAL_MEMORY_HANDLE_SUPPORTED_TYPE    VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT_KHR
#endif

namespace Khaos
{
    static uint32_t _findMemoryType(VkPhysicalDevice phyDevice, uint32_t typeFilter, VkMemoryPropertyFlags properties)
    {
        VkPhysicalDeviceMemoryProperties memProperties;
        vkGetPhysicalDeviceMemoryProperties(phyDevice, &memProperties);

        for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++)
        {
            if ((typeFilter & (1 << i)) && (memProperties.memoryTypes[i].propertyFlags & properties) == properties)
            {
                return i;
            }
        }

        return ~0U;
    }

    VkResult vkCreateExternalImage(VkPhysicalDevice phyDevice, VkDevice device, const VkImageCreateInfo* pCreateInfo, VkImage* pImage, VulkanExternalMemoryHandle* pHandle)
    {
        *pHandle = {};

        // create image
        VkImageCreateInfo imageInfo = *pCreateInfo;
        VkExternalMemoryImageCreateInfoKHR extImageCreateInfo = { VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_IMAGE_CREATE_INFO_KHR };
        extImageCreateInfo.handleTypes = EXTERNAL_MEMORY_HANDLE_SUPPORTED_TYPE;
        imageInfo.pNext = &extImageCreateInfo;

        VkResult res = vkCreateImage(device, &imageInfo, nullptr, pImage);
        if (res != VK_SUCCESS)
            return res;

        // create memory
        VkMemoryRequirements memReqs;
        vkGetImageMemoryRequirements(device, *pImage, &memReqs);
        pHandle->m_allocationSize = memReqs.size;

        VkMemoryAllocateInfo memAlloc = { VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO };
        memAlloc.allocationSize = memReqs.size;
        memAlloc.memoryTypeIndex = _findMemoryType(phyDevice, memReqs.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

        VkExportMemoryAllocateInfoKHR exportInfo = { VK_STRUCTURE_TYPE_EXPORT_MEMORY_ALLOCATE_INFO_KHR };
        exportInfo.handleTypes = EXTERNAL_MEMORY_HANDLE_SUPPORTED_TYPE;
        memAlloc.pNext = &exportInfo;

        res = vkAllocateMemory(device, &memAlloc, nullptr, &pHandle->m_memory);
        if (res != VK_SUCCESS)
            return res;

        // bind memory to image
        res = vkBindImageMemory(device, *pImage, pHandle->m_memory, 0);
        return res;
    }

    void vkDestroyExternalImage(VkDevice device, VkImage image, VulkanExternalMemoryHandle handle)
    {
        vkDestroyImage(device, image, nullptr);
        vkFreeMemory(device, handle.m_memory, nullptr);
    }

    void* getExportHandle(VkDevice device, VulkanExternalMemoryHandle hExtMem)
    {
        if (!hExtMem)
            return nullptr;

#ifndef _WIN32
        VkMemoryGetFdInfoKHR fdInfo = { VK_STRUCTURE_TYPE_MEMORY_GET_FD_INFO_KHR };
        fdInfo.memory = hExtMem.m_memory;
        fdInfo.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD_BIT_KHR;

        int fd = -1;
        auto vkGetMemoryFdKHR = (PFN_vkGetMemoryFdKHR)vkGetDeviceProcAddr(m_device, "vkGetMemoryFdKHR");
        if (vkGetMemoryFdKHR(device, &fdInfo, &fd) != VK_SUCCESS)
            return nullptr;

        return (void*)fd;
#else
        VkMemoryGetWin32HandleInfoKHR handleInfo = { VK_STRUCTURE_TYPE_MEMORY_GET_WIN32_HANDLE_INFO_KHR };
        handleInfo.memory = hExtMem.m_memory;
        handleInfo.handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_WIN32_BIT_KHR;

        HANDLE handle;
        auto vkGetMemoryWin32HandleKHR = (PFN_vkGetMemoryWin32HandleKHR)vkGetDeviceProcAddr(device, "vkGetMemoryWin32HandleKHR");
        if (vkGetMemoryWin32HandleKHR(device, &handleInfo, &handle) != VK_SUCCESS)
            return nullptr;

        return (void*)handle;
#endif
    }
}

